package com.tianji.aigc.controller;

import cn.hutool.core.collection.CollStreamUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@Slf4j
@RestController
@RequiredArgsConstructor
@RequestMapping("/embedding")
public class EmbeddingController {

    private final VectorStore vectorStore;
    private final EmbeddingModel embeddingModel;

    /**
     * 保存数据到向量数据库
     * @param messages 数据列表
     */
    @PostMapping
    public void saveVectorStore(@RequestParam("messages")List<String> messages){
        log.info("将数据保存到向量数据库中{}", messages);
        List<Document> documents = CollStreamUtil.toList(messages, message -> Document.builder()
                .text(message)
                .build());
        this.vectorStore.add(documents);
        log.info("保存到向量数据库成功，数量{}", messages.size());
    }

    /**
     * 将文本向量化
     * @param message 文本
     * @return 向量
     */
    @GetMapping
    public EmbeddingResponse embed(@RequestParam("message") String message){
        log.info("将文本向量化，文本{}", message);
        return embeddingModel.embedForResponse(List.of(message));
    }

    /**
     * 搜索向量数据库
     * @param message 搜索文本
     * @return 搜索结果
     */
    @GetMapping("/search")
    public List<Document> search(@RequestParam("message") String message){
        log.info("搜索向量数据库，文本{}", message);
        return vectorStore.similaritySearch(SearchRequest.builder()
                .query(message).similarityThreshold(0.6d).topK(6).build());
    }

    /**
     * 搜索向量数据库，返回所有数据
     * @return 所有数据
     */
    @GetMapping("/search/all")
    public List<Document> searchAll(){
        log.info("搜索向量数据库，所有数据");
        return vectorStore.similaritySearch(SearchRequest.builder().query("").topK(1000).build());
    }

    /**
     * 根据id批量删除向量数据库中的数据
     * @param ids ids列表
     */
    @DeleteMapping()
    public void delete(@RequestParam("ids") List<String> ids){
        log.info("删除向量数据库，ids{}", ids);
        vectorStore.delete(ids);
    }

}
